
import torch

import altair as alt
import pandas as pd

# LambdaLR 用于自定义 学习率 调整
from torch.optim.lr_scheduler import LambdaLR

# 学习率 使用了 warm_up 策略，因为学习率最开始不能太大，容易找到错误的方向进行训练，
# 所以需要一个 warm_up 的方式让 学习率从小值开始向上升，
# 当 step_num = warmup_steps 的时候， step_num ^ {-0.5} = step_num * warm_up_steps ^ {-1.5}
# 然后 学习率开始 按照指数的方式进行下降
# 经历一次 前向传播 + 反向传播 的 过程，就叫 一个 step，通常是一个 batch
# $$
# lrate = d_{\text{model}}^{-0.5} \cdot
#   \min({step\_num}^{-0.5},
#     {step\_num} \cdot {warmup\_steps}^{-1.5})
# $$

# 是 transformer 论文中的学习率设置，
# 学习率 使用了 warm_up 策略，因为学习率最开始不能太大，容易找到错误的方向进行训练，
# 所以需要一个 warm_up 的方式让 学习率从小值开始向上升，
# 当 step_num = warmup_steps 的时候， step_num ^ {-0.5} = step_num * warm_up_steps ^ {-1.5}
# 然后 学习率开始 按照指数的方式进行下降
# 经历一次 前向传播 + 反向传播 的 过程，就叫 一个 step，通常是一个 batch
# $$
# lrate = d_{\text{model}}^{-0.5} \cdot
#   \min({step\_num}^{-0.5},
#     {step\_num} \cdot {warmup\_steps}^{-1.5})
# $$
# @param (step) : 当前是训练的 第几个 step
# @param (model_size) : d_model
# @param (factor) : 学习率的整体缩放因子
# @param (warmup) : 前多少步使用 warmup
def learningRate(step, model_size : int, factor : float, warmup: int):
    if step == 0:  # 防止被除数是 0
        step = 1
    return factor * (
        model_size ** (-0.5) * min(step ** (-0.5), step * warmup ** (-1.5))
    )


def newWarmupOptim(
        model,
        model_size:int=512,
        base_lr:float=1.0,
        warmup:int=2000,
        factor:float=1.0
):
    optimizer = torch.optim.Adam(
        model.parameters(), lr=base_lr, betas=(0.9, 0.98), eps=1e-9
    )
    # learning rate warm up
    lr_scheduler = LambdaLR(
        optimizer=optimizer,
        lr_lambda=lambda step: learningRate(
            step, model_size, factor=factor, warmup=warmup
        ),
    )
    return optimizer, lr_scheduler


if __name__ == "__main__":

    # 定义三个 example -> d_model, factor, warm_up
    opts = [
        [512, 1, 4000],  # example 1
        [512, 1, 8000],  # example 2
        [256, 1, 4000],  # example 3
    ]

    # 创建一个 简单的示意网络
    dummy_model = torch.nn.Linear(1, 1)
    # 记录每一个 steps 的 learning rate
    learning_rates = []

    # we have 3 examples in opts list.
    for idx, example in enumerate(opts):
        # run 20000 epoch for each example
        optimizer = torch.optim.Adam(
            dummy_model.parameters(), lr=1, betas=(0.9, 0.98), eps=1e-9
        )
        # 设置学习略的更新策略
        lr_scheduler = LambdaLR(
            optimizer=optimizer, lr_lambda=lambda step: learningRate(step, *example)
        )
        tmp = []
        # 一共 20k 的 step
        for step in range(20000):
            tmp.append(optimizer.param_groups[0]["lr"])
            optimizer.step()
            lr_scheduler.step()
        learning_rates.append(tmp)

    learning_rates = torch.tensor(learning_rates)

    # Enable altair to handle more than 5000 rows
    alt.data_transformers.disable_max_rows()

    # 拼接数据
    opts_data = pd.concat(
        [
            pd.DataFrame(
                {
                    "Learning Rate": learning_rates[warmup_idx, :],
                    "model_size:warmup": ["512:4000", "512:8000", "256:4000"][
                        warmup_idx
                    ],
                    "step": range(20000),
                }
            )
            for warmup_idx in [0, 1, 2]
        ]
    )

    # 绘图
    chart = (
        alt.Chart(opts_data)
        .mark_line()
        .properties(width=600)
        .encode(x="step", y="Learning Rate", color="model_size\\:warmup:N")
        .interactive()
    )

    # 保存为 HTML
    chart.save('learning_rate_chart.html')
